import sys
import os

cur_path = os.path.split(__file__)[0]
sys.path.append(cur_path)

import torch
import torch.nn as nn

from BuildingBlocks import (
    Encoder,
    Decoder,
    FinalConv,
    DoubleConv,
    ExtResNetBlock,
    SingleConv,
)


def create_feature_maps(init_channel_number, number_of_fmaps):
    return [init_channel_number * 2**k for k in range(number_of_fmaps)]


class UNet3D(nn.Module):
    """
    3DUnet model from
    `"3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation"
        <https://arxiv.org/pdf/1606.06650.pdf>`.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output segmentation masks;
            Note that that the of out_channels might correspond to either
            different semantic classes or to different binary segmentation mask.
            It's up to the user of the class to interpret the out_channels and
            use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
            or BCEWithLogitsLoss (two-class) respectively)
        f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
            of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
        final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the
            final 1x1 convolution, otherwise apply nn.Softmax. MUST be True if nn.BCELoss (two-class) is used
            to train the model. MUST be False if nn.CrossEntropyLoss (multi-class) is used to train the model.
        layer_order (string): determines the order of layers
            in `SingleConv` module. e.g. 'crg' stands for Conv3d+ReLU+GroupNorm3d.
            See `SingleConv` for more info
        init_channel_number (int): number of feature maps in the first conv layer of the encoder; default: 64
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(
        self,
        in_channels,
        nlayers=4,
        f_maps=16,
        layer_order="crg",
        num_groups=8,
        **kwargs
    ):
        super(UNet3D, self).__init__()

        if isinstance(f_maps, int):
            # use 4 levels in the encoder path as suggested in the paper
            f_maps = create_feature_maps(f_maps, number_of_fmaps=nlayers)

        # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
        # uses DoubleConv as a basic_module for the Encoder
        encoders = []
        for i, out_feature_num in enumerate(f_maps):
            if i == 0:
                encoder = Encoder(
                    in_channels,
                    out_feature_num,
                    apply_pooling=False,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            else:
                encoder = Encoder(
                    f_maps[i - 1],
                    out_feature_num,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            encoders.append(encoder)

        self.encoders = nn.ModuleList(encoders)

        # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
        # uses DoubleConv as a basic_module for the Decoder
        decoders = []
        reversed_f_maps = list(reversed(f_maps))
        for i in range(len(reversed_f_maps) - 1):
            in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
            out_feature_num = reversed_f_maps[i + 1]
            decoder = Decoder(
                in_feature_num,
                out_feature_num,
                basic_module=DoubleConv,
                conv_layer_order=layer_order,
                num_groups=num_groups,
            )
            decoders.append(decoder)

        self.decoders = nn.ModuleList(decoders)

    def forward(self, x):
        # encoder part
        encoders_features = []
        for encoder in self.encoders:
            x = encoder(x)
            # reverse the encoder outputs to be aligned with the decoder
            encoders_features.insert(0, x)

        # remove the last encoder's output from the list
        # !!remember: it's the 1st in the list

        encoders_features = encoders_features[1:]
        # decoder part
        for decoder, encoder_features in zip(self.decoders, encoders_features):
            # pass the output from the corresponding encoder and the output
            # of the previous decoder
            x = decoder(encoder_features, x)

        # apply final_activation (i.e. Sigmoid or Softmax) only for prediction. During training the network outputs
        # logits and it's up to the user to normalize it before visualising with tensorboard or computing validation metric

        return x


class UNet3DWithSamplePoints(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        num_levels: int,
        use_position_input: bool = True,
    ):
        super().__init__()
        self.unet3d = UNet3D(
            in_channels=in_channels, nlayers=num_levels, f_maps=hidden_channels
        )
        self.hidden_channels = hidden_channels
        self.out_ln = nn.LayerNorm(hidden_channels)

    def voxel_expand(self, voxel_left):  # B C Z X Y
        indices = torch.arange(31, -1, -1, dtype=torch.long, device=voxel_left.device)
        voxel_right = voxel_left[:, :, :, indices, :]
        voxel_merge = torch.cat((voxel_left, voxel_right), axis=3)
        return voxel_merge

    def forward(
        self,
        x,
        output_points,
        graph_node=None,
        graph_cell=None,
        half=False,
        is_training=True,
    ):

        # Use super class to get the output of the UNet3D
        x = self.unet3d.forward(x)
        # sample_points to 5D by adding two dimensions
        if half:
            x = self.voxel_expand(x)

        """ interpolate the output of the UNet3D to the node_wise points"""
        node_x = torch.nn.functional.grid_sample(
            x, graph_node.query, align_corners=False
        )
        # Remove the two dimensions
        node_x = node_x.squeeze(3).squeeze(3)
        # (BxCxN) -> (BxNxC)
        node_x = node_x.permute(0, 2, 1).reshape(-1, self.hidden_channels)

        return self.out_ln(node_x)


if __name__ == "__main__":
    net = UNet3DWithSamplePoints(1, 64, 64, 4, True)
    sdf = torch.randn([7, 1, 64, 64, 64])
    outputpoints = torch.randn([7, 1471, 1, 1, 3])
    output = net.forward(sdf, outputpoints, half=False)
    print(output.shape)
